#include "../Base.h"

#include <ntimage.h>

typedef struct _IMAGE_RELOC{

	UINT16 Offset : 12; //偏移
	UINT16 Type : 4;  //类型   3:32位 ,10:64位
}IMAGE_RELOC, * PIMAGE_RELOC;

typedef NTSTATUS(NTAPI* DriverEntryProc)(PDRIVER_OBJECT pDriver, PUNICODE_STRING pReg);

char* FileToImage(char* fileBuffer) {
	PIMAGE_DOS_HEADER pDos = (PIMAGE_DOS_HEADER)fileBuffer;
	PIMAGE_NT_HEADERS pNt = (PIMAGE_NT_HEADERS)(fileBuffer + pDos->e_lfanew);
	ULONG sizeofImage = pNt->OptionalHeader.SizeOfImage;
	
	//创建imageBuffer
	char* imageBuffer = (char*)ExAllocatePoolWithTag(NonPagedPool, sizeofImage, 'yang');
	memset(imageBuffer, 0, sizeofImage);


	//复制PE头
	memcpy(imageBuffer, fileBuffer, pNt->OptionalHeader.SizeOfHeaders);

	ULONG NumberOfSections = pNt->FileHeader.NumberOfSections;


	PIMAGE_SECTION_HEADER pSection = IMAGE_FIRST_SECTION(pNt);

	for (ULONG i = 0; i < NumberOfSections; i++) {
		
		memcpy(imageBuffer + pSection->VirtualAddress, fileBuffer + pSection->PointerToRawData, pSection->SizeOfRawData);
		pSection++;

	}
	return imageBuffer;
}

BOOLEAN UpdataReloc(char* imageBuffer) {

	PIMAGE_DOS_HEADER pDos = (PIMAGE_DOS_HEADER)imageBuffer;
	PIMAGE_NT_HEADERS pNt = (PIMAGE_NT_HEADERS)(imageBuffer + pDos->e_lfanew);

	if (!pNt) return FALSE;
	
	PIMAGE_DATA_DIRECTORY iRelocation = &(IMAGE_DATA_DIRECTORY)(pNt->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_BASERELOC]);
	PIMAGE_BASE_RELOCATION pBase = (PIMAGE_BASE_RELOCATION)(imageBuffer + iRelocation->VirtualAddress);

	while (pBase->SizeOfBlock && pBase->VirtualAddress) {

		PIMAGE_RELOC RelocationBlock = (PIMAGE_RELOC)((PUCHAR)pBase + sizeof(IMAGE_BASE_RELOCATION));
		UINT32 NumberOfRelocations = (pBase->SizeOfBlock - sizeof(IMAGE_BASE_RELOCATION)) / sizeof(IMAGE_RELOC);
		for (int i = 0; i < NumberOfRelocations; i++) {

			
			if (RelocationBlock[i].Type == IMAGE_REL_BASED_DIR64) {
				//64位
				PUINT64 address = (PUINT64)((PUINT8)imageBuffer + pBase->VirtualAddress + RelocationBlock[i].Offset);
				UINT64 Delta =(UINT64)(*address - pNt->OptionalHeader.ImageBase + (UINT64)imageBuffer);  
				*address = Delta;
			}
			else if (RelocationBlock[i].Type == IMAGE_REL_BASED_HIGHLOW) {

				PUINT32 address = (PUINT32)((PUINT8)imageBuffer + pBase->VirtualAddress + RelocationBlock[i].Offset);
				UINT32 Delta = (UINT32)(*address - pNt->OptionalHeader.ImageBase + (UINT32)imageBuffer);
				*address = Delta;
			}
		}
		pBase = (PIMAGE_BASE_RELOCATION)(PUCHAR)pBase + pBase->SizeOfBlock;
	}
	return TRUE;
}


BOOLEAN UpdateImportTable(char* imageBuffer) {

	PIMAGE_DOS_HEADER pDos = (PIMAGE_DOS_HEADER)imageBuffer;
	PIMAGE_NT_HEADERS pNt = (PIMAGE_NT_HEADERS)(imageBuffer + pDos->e_lfanew);

	if (!pNt) return FALSE;

	PIMAGE_DATA_DIRECTORY pd = &(IMAGE_DATA_DIRECTORY)(pNt->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_IMPORT]);
	PIMAGE_IMPORT_DESCRIPTOR pI = (PIMAGE_IMPORT_DESCRIPTOR)(imageBuffer + pd->VirtualAddress);
	
	BOOLEAN isSuccess = TRUE;

	for (; pI->Name; pI++) {

		char* moduleName  = imageBuffer + pI->Name;
		ULONG_PTR base = QuerySystemModuleInfo(moduleName, NULL);
		if (!base) {

			isSuccess = FALSE;
			break;
		}
		PIMAGE_THUNK_DATA pThuckName = (PIMAGE_THUNK_DATA)(imageBuffer + pI->OriginalFirstThunk); //第一次加载到内存时保存名字的地方
		PIMAGE_THUNK_DATA pThunkFunc = (PIMAGE_THUNK_DATA)(imageBuffer + pI->FirstThunk);  //后续调用使用的地方


		for (; pThuckName->u1.ForwarderString; ++pThunkFunc, ++pThuckName) {
			PIMAGE_IMPORT_BY_NAME FuncName = (PIMAGE_IMPORT_BY_NAME)(imageBuffer + pThuckName->u1.AddressOfData);
			
			ULONG_PTR func = (ULONG_PTR)RtlFindExportedRoutineByName((char*)base, FuncName->Name);
			if (func) {

				pThunkFunc->u1.Function = func;
			}
			else {
				isSuccess = FALSE;
				break;
			}


		}

		if (!isSuccess) break;
	}

}


//��ȡ�� LoadLibraryExW
ULONG64 ExportTableFuncByName(char* pData, char* funcName)
{
	PIMAGE_DOS_HEADER pHead = (PIMAGE_DOS_HEADER)pData;
	PIMAGE_NT_HEADERS pNt = (PIMAGE_NT_HEADERS)(pData + pHead->e_lfanew);
	int numberRvaAndSize = pNt->OptionalHeader.NumberOfRvaAndSizes;
	PIMAGE_DATA_DIRECTORY pDir = (PIMAGE_DATA_DIRECTORY)&pNt->OptionalHeader.DataDirectory[0];
	
	PIMAGE_EXPORT_DIRECTORY pExport = (PIMAGE_EXPORT_DIRECTORY)(pData + pDir->VirtualAddress);
	ULONG64 funcAddr = 0;
	for (int i = 0; i < pExport->NumberOfNames; i++)
	{
		int* funcAddress = (int*)(pData + pExport->AddressOfFunctions);//乱序
		int* names = (int*)(pData + pExport->AddressOfNames); //有序
		short* fh = (short*)(pData + pExport->AddressOfNameOrdinals);//有序
		int index = -1;
		char* name = pData + names[i];

		if (strcmp(name, funcName) == 0)
		{
			index = fh[i];
		}



		if (index != -1)
		{

			funcAddr = ULONG(pData + funcAddress[index]);
			break;
		}


	}

	if (!funcAddr)
	{
		KdPrint(("û���ҵ�����%s\r\n", funcName));
		return NULL;

	}



	return funcAddr;
}




#define DEREF( name )*(UINT_PTR *)(name)
#define DEREF_64( name )*(unsigned __int64 *)(name)
#define DEREF_32( name )*(unsigned long *)(name)
#define DEREF_16( name )*(unsigned short *)(name)
#define DEREF_8( name )*(UCHAR *)(name)

ULONG_PTR ExportProcAddressByModuleBase(ULONG_PTR hModule, const char* lpProcName, BOOLEAN x64Module)
{
	UINT_PTR uiLibraryAddress = 0;
	ULONG_PTR fpResult = NULL;

	if (hModule == NULL)
		return NULL;

	// a module handle is really its base address
	uiLibraryAddress = (UINT_PTR)hModule;

	__try
	{
		UINT_PTR uiAddressArray = 0;
		UINT_PTR uiNameArray = 0;
		UINT_PTR uiNameOrdinals = 0;
		PIMAGE_NT_HEADERS32 pNtHeaders32 = NULL;
		PIMAGE_NT_HEADERS64 pNtHeaders64 = NULL;
		PIMAGE_DATA_DIRECTORY pDataDirectory = NULL;
		PIMAGE_EXPORT_DIRECTORY pExportDirectory = NULL;

		// get the VA of the modules NT Header
		pNtHeaders32 = (PIMAGE_NT_HEADERS32)(uiLibraryAddress + ((PIMAGE_DOS_HEADER)uiLibraryAddress)->e_lfanew);
		pNtHeaders64 = (PIMAGE_NT_HEADERS64)(uiLibraryAddress + ((PIMAGE_DOS_HEADER)uiLibraryAddress)->e_lfanew);
		if (x64Module)
		{
			pDataDirectory = (PIMAGE_DATA_DIRECTORY)&pNtHeaders64->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_EXPORT];
		}
		else
		{
			pDataDirectory = (PIMAGE_DATA_DIRECTORY)&pNtHeaders32->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_EXPORT];
		}


		// get the VA of the export directory
		pExportDirectory = (PIMAGE_EXPORT_DIRECTORY)(uiLibraryAddress + pDataDirectory->VirtualAddress);

		// get the VA for the array of addresses
		uiAddressArray = (uiLibraryAddress + pExportDirectory->AddressOfFunctions);

		// get the VA for the array of name pointers
		uiNameArray = (uiLibraryAddress + pExportDirectory->AddressOfNames);

		// get the VA for the array of name ordinals
		uiNameOrdinals = (uiLibraryAddress + pExportDirectory->AddressOfNameOrdinals);

		// test if we are importing by name or by ordinal...
		if ((PtrToUlong(lpProcName) & 0xFFFF0000) == 0x00000000)
		{
			// import by ordinal...

			// use the import ordinal (- export ordinal base) as an index into the array of addresses
			uiAddressArray += ((IMAGE_ORDINAL(PtrToUlong(lpProcName)) - pExportDirectory->Base) * sizeof(unsigned long));

			// resolve the address for this imported function
			fpResult = (ULONG_PTR)(uiLibraryAddress + DEREF_32(uiAddressArray));
		}
		else
		{
			// import by name...
			unsigned long dwCounter = pExportDirectory->NumberOfNames;
			while (dwCounter--)
			{
				char* cpExportedFunctionName = (char*)(uiLibraryAddress + DEREF_32(uiNameArray));

				// test if we have a match...
				if (strcmp(cpExportedFunctionName, lpProcName) == 0)
				{
					// use the functions name ordinal as an index into the array of name pointers
					uiAddressArray += (DEREF_16(uiNameOrdinals) * sizeof(unsigned long));

					// calculate the virtual address for the function
					fpResult = (ULONG_PTR)(uiLibraryAddress + DEREF_32(uiAddressArray));

					// finish...
					break;
				}

				// get the next exported function name
				uiNameArray += sizeof(unsigned long);

				// get the next exported function name ordinal
				uiNameOrdinals += sizeof(unsigned short);
			}
		}
	}
	__except (EXCEPTION_EXECUTE_HANDLER)
	{
		fpResult = NULL;
	}

	return fpResult;
}



BOOLEAN UpdateCookie(char* imageBuffer) {

	PIMAGE_NT_HEADERS pNts = (PIMAGE_NT_HEADERS)RtlImageNtHeader(imageBuffer);
	if (!pNts) return FALSE;
	PIMAGE_DATA_DIRECTORY pConfigDir = &pNts->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_LOAD_CONFIG];
	PIMAGE_LOAD_CONFIG_DIRECTORY pConfig = (PIMAGE_LOAD_CONFIG_DIRECTORY)(imageBuffer + pConfigDir->VirtualAddress);
	*(PULONG_PTR)(pConfig->SecurityCookie) += 10;

}

//必须关闭VS的CFG control flow  guard
BOOLEAN LoadDriver(PUCHAR fileBuffer) {

	char* imageBase = FileToImage((char*)fileBuffer);
	if (!imageBase) {
		return FALSE;
	}
	BOOLEAN isSuccess = FALSE;

	do {
		isSuccess = UpdataReloc(imageBase);
		if (!isSuccess) {
			break;
		}
		isSuccess = UpdateImportTable(imageBase);
		if (!isSuccess) {
			break;
		}

		//修复Cookie
		UpdateCookie(imageBase);
		//call 入口点
		PIMAGE_NT_HEADERS pNts = (PIMAGE_NT_HEADERS)RtlImageNtHeader(imageBase);
		ULONG_PTR entry = pNts->OptionalHeader.AddressOfEntryPoint;
		DriverEntryProc dirverEntry = (DriverEntryProc)(imageBase + entry);
		NTSTATUS status = dirverEntry(NULL, NULL);
		if (NT_SUCCESS(status)) {
			isSuccess = FALSE;
			break;
		}
		//清空PE头
		memset(imageBase, 0, PAGE_SIZE);

	} while (0);

	if (!isSuccess) {
		ExFreePoolWithTag(imageBase, 'yang');
		return FALSE;
	}



	ExFreePoolWithTag(imageBase, 'yang');

	return isSuccess;

}


BOOLEAN LoadDriverTeset() {
	char sysData[] = "123123";
	ULONG dwImageSize = sizeof(sysData);
	UCHAR* pMemory = (UCHAR*)ExAllocatePoolWithTag(NonPagedPool, dwImageSize, 'yang');
	memcpy(pMemory, sysData, dwImageSize);
	LoadDriver(pMemory);
	ExFreePoolWithTag((PVOID)pMemory, 'yang');
	return TRUE;
}